package document

import (
	"context"
	"fmt"
	"os"
	"path/filepath"

	"github.com/cloudwego/eino-ext/components/document/transformer/splitter/markdown"
	"github.com/cloudwego/eino-ext/components/document/transformer/splitter/recursive"
	"github.com/cloudwego/eino-ext/components/document/transformer/splitter/semantic"
	"github.com/cloudwego/eino/components/embedding"
	"github.com/cloudwego/eino/schema"
	"github.com/pelletier/go-toml/v2"
)

// 导入errors.go中定义的错误变量

// DocumentSplitter 定义文档分割器接口
type DocumentSplitter interface {
	Transform(ctx context.Context, docs []*schema.Document) ([]*schema.Document, error)
}

// SplitterConfig TOML配置文件结构
type SplitterConfig struct {
	RecursiveSplitter struct {
		ChunkSize   int      `toml:"chunk_size"`
		OverlapSize int      `toml:"overlap_size"`
		Separators  []string `toml:"separators"`
		KeepType    string   `toml:"keep_type"`
	} `toml:"recursive_splitter"`
	SemanticSplitter struct {
		BufferSize   int      `toml:"buffer_size"`
		MinChunkSize int      `toml:"min_chunk_size"`
		Separators   []string `toml:"separators"`
		Percentile   float64  `toml:"percentile"`
	} `toml:"semantic_splitter"`
	MarkdownSplitter struct {
		HeadingPriority []string `toml:"heading_priority"`
		MaxChunkSize    int      `toml:"max_chunk_size"`
		OverlapSize     int      `toml:"overlap_size"`
		Headers         struct {
			H1 string `toml:"#"`
			H2 string `toml:"##"`
			H3 string `toml:"###"`
		} `toml:"headers"`
		TrimHeaders bool `toml:"trim_headers"`
	} `toml:"markdown_splitter"`
}

// RecursiveSplitterConfig 递归分割器配置
type RecursiveSplitterConfig struct {
	ChunkSize   int      // 块大小（字符数）
	OverlapSize int      // 重叠大小（字符数）
	Separators  []string // 分隔符列表
	KeepType    string   // 保留类型（"start", "end"）
}

// LoadConfig 从TOML文件加载分割器配置
func LoadConfig(configPath string) (*SplitterConfig, error) {
	// 读取配置文件
	data, err := os.ReadFile(configPath)
	if err != nil {
		return nil, fmt.Errorf("读取配置文件失败: %w", err)
	}

	// 解析TOML配置
	var config SplitterConfig
	err = toml.Unmarshal(data, &config)
	if err != nil {
		return nil, fmt.Errorf("解析TOML配置失败: %w", err)
	}

	return &config, nil
}

// NewRecursiveSplitter 创建递归分割器并执行分割
// 参数:
//   - ctx: 上下文
//   - docs: 需要分割的文档列表（从请求参数传入）
//
// 返回:
//   - []*schema.Document: 分割后的文档片段列表
//   - error: 错误信息
func NewRecursiveSplitter(ctx context.Context, docs []*schema.Document) ([]*schema.Document, error) {
	// 参数验证
	if len(docs) == 0 {
		return nil, ErrEmptyDocumentList
	}

	// 获取当前文件所在目录
	dir, err := filepath.Abs(filepath.Dir(os.Args[0]))
	if err != nil {
		dir = "."
	}

	// 配置文件路径
	configPath := filepath.Join(dir, "splitter_config.toml")

	// 加载TOML配置
	tomlConfig, err := LoadConfig(configPath)
	if err != nil {
		// 如果配置文件不存在，使用默认配置
		fmt.Printf("警告: 无法加载配置文件 %s，使用默认配置: %v\n", configPath, err)
		tomlConfig = &SplitterConfig{}
		tomlConfig.RecursiveSplitter.ChunkSize = 1000
		tomlConfig.RecursiveSplitter.OverlapSize = 200
		tomlConfig.RecursiveSplitter.Separators = []string{"\n\n", "\n", "。", "！", "？"}
		tomlConfig.RecursiveSplitter.KeepType = "end"
	}

	// 初始化配置
	recursiveConfig := &recursive.Config{
		ChunkSize:   tomlConfig.RecursiveSplitter.ChunkSize,
		OverlapSize: tomlConfig.RecursiveSplitter.OverlapSize,
		Separators:  tomlConfig.RecursiveSplitter.Separators,
		KeepType:    recursive.KeepTypeEnd,
	}

	// 设置KeepType
	if tomlConfig.RecursiveSplitter.KeepType == "start" {
		recursiveConfig.KeepType = recursive.KeepTypeStart
	} else if tomlConfig.RecursiveSplitter.KeepType == "end" {
		recursiveConfig.KeepType = recursive.KeepTypeEnd
	}

	// 配置已从TOML文件自动加载，无需额外配置参数

	// 初始化分割器
	splitter, err := recursive.NewSplitter(ctx, recursiveConfig)
	if err != nil {
		return nil, fmt.Errorf("初始化递归分割器失败: %w", err)
	}
	// 执行分割
	results, err := splitter.Transform(ctx, docs)
	if err != nil {
		panic(err)
	}

	return results, nil
}

// SplitDocumentRecursive 使用递归分割器分割文档
// 参数:
//   - ctx: 上下文
//   - docs: 要分割的文档列表（从请求参数传入）
//
// 返回:
//   - []*schema.Document: 分割后的文档片段列表
//   - error: 错误信息
func SplitDocumentRecursive(ctx context.Context, docs []*schema.Document) ([]*schema.Document, error) {
	// 参数验证
	if len(docs) == 0 {
		return nil, ErrEmptyDocumentList
	}

	// 直接调用NewRecursiveSplitter进行分割，配置将自动初始化
	results, err := NewRecursiveSplitter(ctx, docs)
	if err != nil {
		return nil, fmt.Errorf("递归分割文档失败: %w", err)
	}

	return results, nil
}

// SemanticSplitterConfig 语义分割器配置
type SemanticSplitterConfig struct {
	BufferSize   int      // 缓冲区大小
	MinChunkSize int      // 最小块大小
	Separators   []string // 分隔符列表
	Percentile   float64  // 百分位阈值
}

// NewSemanticSplitter 创建语义分割器并执行分割
// 参数:
//   - ctx: 上下文
//   - docs: 需要分割的文档列表（从请求参数传入）
//
// 返回:
//   - []*schema.Document: 分割后的文档片段列表
//   - error: 错误信息
func NewSemanticSplitter(ctx context.Context, docs []*schema.Document) ([]*schema.Document, error) {
	// 参数验证
	if len(docs) == 0 {
		return nil, ErrEmptyDocumentList
	}

	// 获取当前文件所在目录
	dir, err := filepath.Abs(filepath.Dir(os.Args[0]))
	if err != nil {
		dir = "."
	}

	// 配置文件路径
	configPath := filepath.Join(dir, "splitter_config.toml")

	// 加载TOML配置
	tomlConfig, err := LoadConfig(configPath)
	if err != nil {
		// 如果配置文件不存在，使用默认配置
		fmt.Printf("警告: 无法加载配置文件 %s，使用默认配置: %v\n", configPath, err)
		tomlConfig = &SplitterConfig{}
		tomlConfig.SemanticSplitter.BufferSize = 5
		tomlConfig.SemanticSplitter.MinChunkSize = 300
		tomlConfig.SemanticSplitter.Separators = []string{"\n\n", "\n", "。", "！", "？"}
		tomlConfig.SemanticSplitter.Percentile = 0.9
	}

	// 语义分割器需要外部注入 Embedding 实现，此处仅作占位
	// 实际使用时需传入具体实现（如 openai、baichuan 等）
	var embedder embedding.Embedder

	// 初始化分割器
	splitter, err := semantic.NewSplitter(ctx, &semantic.Config{
		Embedding:    embedder,
		BufferSize:   tomlConfig.SemanticSplitter.BufferSize,
		MinChunkSize: tomlConfig.SemanticSplitter.MinChunkSize,
		Separators:   tomlConfig.SemanticSplitter.Separators,
		Percentile:   tomlConfig.SemanticSplitter.Percentile,
	})
	if err != nil {
		panic(err)
	}

	// 执行分割
	results, err := splitter.Transform(ctx, docs)
	if err != nil {
		panic(err)
	}
	return results, nil
}

// SplitDocumentSemantic 使用语义分割器分割文档
// 参数:
//   - ctx: 上下文
//   - docs: 要分割的文档列表（从请求参数传入）
//
// 返回:
//   - []*schema.Document: 分割后的文档片段列表
//   - error: 错误信息
func SplitDocumentSemantic(ctx context.Context, docs []*schema.Document) ([]*schema.Document, error) {
	// 参数验证
	if len(docs) == 0 {
		return nil, ErrEmptyDocumentList
	}

	// 直接调用NewSemanticSplitter进行分割，配置将自动初始化
	results, err := NewSemanticSplitter(ctx, docs)
	if err != nil {
		return nil, fmt.Errorf("语义分割文档失败: %w", err)
	}

	return results, nil
}

// MarkdownSplitterConfig Markdown分割器配置
type MarkdownSplitterConfig struct {
	HeadingPriority []string          // 标题优先级列表
	MaxChunkSize    int               // 最大块大小
	OverlapSize     int               // 重叠大小
	Headers         map[string]string // 标题映射（如 "#": "h1"）
	TrimHeaders     bool              // 是否修剪标题
}

// NewMarkdownSplitter 创建Markdown分割器并执行分割
// 参数:
//   - ctx: 上下文
//   - docs: 需要分割的文档列表（从请求参数传入）
//
// 返回:
//   - []*schema.Document: 分割后的文档片段列表
//   - error: 错误信息
func NewMarkdownSplitter(ctx context.Context, docs []*schema.Document) ([]*schema.Document, error) {
	// 参数验证
	if len(docs) == 0 {
		return nil, ErrEmptyDocumentList
	}

	// 获取当前文件所在目录
	dir, err := filepath.Abs(filepath.Dir(os.Args[0]))
	if err != nil {
		dir = "."
	}

	// 配置文件路径
	configPath := filepath.Join(dir, "splitter_config.toml")

	// 加载TOML配置
	tomlConfig, err := LoadConfig(configPath)
	if err != nil {
		// 如果配置文件不存在，使用默认配置
		fmt.Printf("警告: 无法加载配置文件 %s，使用默认配置: %v\n", configPath, err)
		tomlConfig = &SplitterConfig{}
		tomlConfig.MarkdownSplitter.HeadingPriority = []string{"#", "##", "###", "####", "#####", "######"}
		tomlConfig.MarkdownSplitter.MaxChunkSize = 1000
		tomlConfig.MarkdownSplitter.OverlapSize = 0
	}

	// 将HeadingPriority转换为Headers映射
	headersMap := make(map[string]string)
	for _, header := range tomlConfig.MarkdownSplitter.HeadingPriority {
		level := len(header) // 标题级别由#的数量决定
		headersMap[header] = fmt.Sprintf("h%d", level)
	}

	// 初始化分割器，使用从TOML加载的配置
	splitter, err := markdown.NewHeaderSplitter(ctx, &markdown.HeaderConfig{
		Headers:     headersMap,
		TrimHeaders: false, // 可考虑添加到配置中
	})
	if err != nil {
		return nil, fmt.Errorf("初始化Markdown分割器失败: %w", err)
	}

	// 执行分割，使用传入的文档参数
	results, err := splitter.Transform(ctx, docs)
	if err != nil {
		return nil, fmt.Errorf("执行Markdown文档分割失败: %w", err)
	}

	// 返回实际的分割结果
	return results, nil
}

// SplitDocumentMarkdown 使用Markdown分割器分割文档
// 参数:
//   - ctx: 上下文
//   - docs: 要分割的文档列表（从请求参数传入）
//
// 返回:
//   - []*schema.Document: 分割后的文档片段列表
//   - error: 错误信息
func SplitDocumentMarkdown(ctx context.Context, docs []*schema.Document) ([]*schema.Document, error) {
	// 参数验证
	if len(docs) == 0 {
		return nil, ErrEmptyDocumentList
	}

	// 直接调用NewMarkdownSplitter进行分割，配置将自动初始化
	results, err := NewMarkdownSplitter(ctx, docs)
	if err != nil {
		return nil, fmt.Errorf("Markdown分割文档失败: %w", err)
	}

	return results, nil
}

// SplitDocument 统一的文档分割入口函数
// 参数:
//   - ctx: 上下文
//   - docs: 要分割的文档列表（从请求参数传入）
//   - splitType: 分割类型 ("recursive", "semantic", "markdown")
//
// 返回:
//   - []*schema.Document: 分割后的文档片段列表
//   - error: 错误信息
func SplitDocument(ctx context.Context, docs []*schema.Document, splitType string) ([]*schema.Document, error) {
	// 参数验证
	if len(docs) == 0 {
		return nil, ErrEmptyDocumentList
	}

	// 根据分割类型选择处理方式
	switch splitType {
	case "recursive":
		return SplitDocumentRecursive(ctx, docs)
	case "semantic":
		return SplitDocumentSemantic(ctx, docs)
	case "markdown":
		return SplitDocumentMarkdown(ctx, docs)
	default:
		return nil, fmt.Errorf("不支持的分割类型: %s", splitType)
	}
}
